import os
import torch
from FlagEmbedding import FlagAutoReranker

def test_reranker_example():
    # 检测是否有MPS可用(Apple Silicon芯片)
    if torch.backends.mps.is_available():
        device = ["mps"]
        print("使用MPS加速(Apple Silicon GPU)")
        use_fp16 = True
    else:
        device = ["cpu"]
        print("使用CPU")
        use_fp16 = False
    
    # 加载模型
    model = FlagAutoReranker.from_finetuned(
        'BAAI/bge-reranker-large',  # 也可以尝试基础模型 'BAAI/bge-reranker-base'
        use_fp16=use_fp16,  # 在MPS上可以使用FP16加速
        batch_size=8,    # 根据您的内存调整
        query_max_length=256,
        max_length=512,
        devices=device,
        cache_dir=os.getenv('HF_HUB_CACHE', None),
    )
    
    # 准备一些测试数据
    pairs = [
        ["什么是人工智能？", "人工智能是计算机科学的一个分支，旨在创建能够执行通常需要人类智能的任务的系统。"],
        ["什么是人工智能？", "苹果是一种常见的水果，有红色、绿色等多种颜色。"],
        ["世界上最高的山峰是什么？", "珠穆朗玛峰是世界上最高的山峰，位于中国和尼泊尔边境。"],
        ["世界上最高的山峰是什么？", "人工智能是计算机科学的一个分支，旨在创建能够执行通常需要人类智能的任务的系统。"]
    ]
    
    print("\n计算相关性分数中...")
    # 计算相关性分数
    scores = model.compute_score(pairs)
    print("原始分数:", scores)
    
    # 归一化分数（范围为0-1）
    normalized_scores = model.compute_score(pairs, normalize=True)
    print("归一化分数:", normalized_scores)
    
    # 解释结果
    for i, (pair, score, norm_score) in enumerate(zip(pairs, scores, normalized_scores)):
        print(f"\n示例 {i+1}:")
        print(f"查询: {pair[0]}")
        print(f"段落: {pair[1]}")
        print(f"原始分数: {score:.4f}")
        print(f"归一化分数: {norm_score:.4f}")
        print(f"相关性: {'高' if norm_score > 0.5 else '低'}")

if __name__ == '__main__':
    print("Mac系统上的BGE Reranker示例")
    test_reranker_example()